#!/usr/bin/env python3
import math
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Point
from std_msgs.msg import Float64
from tf2_ros import TransformListener, Buffer
import matplotlib.pyplot as plt
from collections import deque

# === 像素换算&符号（按你原代码习惯，可自行调整）===
PX_PER_M = 546.5  # 米到像素的尺度
# /result3_topic: x=像素(左右), y=像素(上下), z=米(深度)
# TF映射（示例：基于你之前的注释/尝试）
#  - X轴像素对应：-trans.translation.y * PX_PER_M
#  - Y轴像素对应：-trans.translation.z * PX_PER_M
#  - Z轴“米”对应：  trans.translation.x （直接米值；为了可视化一致，也可乘以 PX_PER_M）
REAL_TO_PX_X_SIGN = -1.0
REAL_TO_PX_Y_SIGN = -1.0
REAL_TO_M_Z_SIGN  = +1.0

# “真实数据”额外可视化放大（让真实曲线更醒目，不影响数据本身）
GAIN_REAL_X = 1.5
GAIN_REAL_Y = 1.5
GAIN_REAL_Z = 1.0   # Z若画成米，通常无需额外放大；如需与像素同量纲可把下面乘以 PX_PER_M 并调增此增益

HISTORY = 400       # 窗口长度
DT = 0.05           # 20Hz 绘图

class PlotNode(Node):
    def __init__(self):
        super().__init__('plot_result3_topic')

        # 数据队列
        self.ticks = deque(maxlen=HISTORY)
        self.idx = 0

        # 原始 result3_topic
        self.raw_x = deque(maxlen=HISTORY)
        self.raw_y = deque(maxlen=HISTORY)
        self.raw_z = deque(maxlen=HISTORY)

        # 真实 TF（换算/放大后）
        self.real_x = deque(maxlen=HISTORY)
        self.real_y = deque(maxlen=HISTORY)
        self.real_z = deque(maxlen=HISTORY)

        # ibvs avg
        self.avg_x = deque(maxlen=HISTORY)
        self.avg_y = deque(maxlen=HISTORY)
        self.avg_z = deque(maxlen=HISTORY)

        # 最近值（若无则用 NaN 填充，避免长度不一致）
        self._last_raw = {'x': math.nan, 'y': math.nan, 'z': math.nan}
        self._last_avg = {'x': math.nan, 'y': math.nan, 'z': math.nan}

        # 订阅
        self.create_subscription(Point,  '/result3_topic', self.cb_raw, 10)
        self.create_subscription(Float64,'/ibvs_x_avg',   self.cb_avg_x, 10)
        self.create_subscription(Float64,'/ibvs_y_avg',   self.cb_avg_y, 10)
        self.create_subscription(Float64,'/ibvs_z_avg',   self.cb_avg_z, 10)

        # TF
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)

        # 定时器
        self.timer = self.create_timer(DT, self.on_timer)

        # Matplotlib 三子图
        plt.ion()
        self.fig, (self.ax_x, self.ax_y, self.ax_z) = plt.subplots(3, 1, figsize=(10, 9), sharex=True)
        # X
        self.lx_raw,  = self.ax_x.plot([], [], label='raw topic.x (px)')
        self.lx_real, = self.ax_x.plot([], [], label='real (px) * gain')
        self.lx_avg,  = self.ax_x.plot([], [], label='ibvs_x_avg (px)')
        self.ax_x.set_ylabel('X (px)')
        self.ax_x.legend(loc='upper right')
        self.ax_x.grid(True, alpha=0.3)

        # Y
        self.ly_raw,  = self.ax_y.plot([], [], label='raw topic.y (px)')
        self.ly_real, = self.ax_y.plot([], [], label='real (px) * gain')
        self.ly_avg,  = self.ax_y.plot([], [], label='ibvs_y_avg (px)')
        self.ax_y.set_ylabel('Y (px)')
        self.ax_y.legend(loc='upper right')
        self.ax_y.grid(True, alpha=0.3)

        # Z
        self.lz_raw,  = self.ax_z.plot([], [], label='raw topic.z (m)')
        self.lz_real, = self.ax_z.plot([], [], label='real (m) * gain')
        self.lz_avg,  = self.ax_z.plot([], [], label='ibvs_z_avg (m)')
        self.ax_z.set_ylabel('Z (m)')
        self.ax_z.set_xlabel('Tick')
        self.ax_z.legend(loc='upper right')
        self.ax_z.grid(True, alpha=0.3)

        self.fig.suptitle('Result3 vs TF(real) vs IBVS avg (X/Y/Z)')

    # 回调：原始
    def cb_raw(self, msg: Point):
        self._last_raw['x'] = float(msg.x)
        self._last_raw['y'] = float(msg.y)
        self._last_raw['z'] = float(msg.z)

    # 回调：avg
    def cb_avg_x(self, msg: Float64):
        self._last_avg['x'] = float(msg.data)

    def cb_avg_y(self, msg: Float64):
        self._last_avg['y'] = float(msg.data)

    def cb_avg_z(self, msg: Float64):
        self._last_avg['z'] = float(msg.data)

    # 每个周期：采样 & 画图
    def on_timer(self):
        # tick 递增
        self.ticks.append(self.idx)
        self.idx += 1

        # --- 1) 原始数据入库（确保每次都 push，避免长度不一致）---
        self.raw_x.append(self._last_raw['x'])
        self.raw_y.append(self._last_raw['y'])
        self.raw_z.append(self._last_raw['z'])

        # --- 2) TF -> real 曲线（异常时写 NaN，一样 push）---
        rx, ry, rz = math.nan, math.nan, math.nan
        try:
            trans = self.tf_buffer.lookup_transform('base_link', 'end', rclpy.time.Time())
            # X/Y 映射为“像素”以便与 raw/avg 对齐
            rx = REAL_TO_PX_X_SIGN * trans.transform.translation.y * PX_PER_M * GAIN_REAL_X * (-1.0)  # 等价合并符号亦可
            ry = REAL_TO_PX_Y_SIGN * trans.transform.translation.z * PX_PER_M * GAIN_REAL_Y * (-1.0)
            # Z 以“米”为主（如需 px，可改乘 PX_PER_M 并调整 GAIN_REAL_Z）
            rz = REAL_TO_M_Z_SIGN  * trans.transform.translation.x * GAIN_REAL_Z
        except Exception:
            pass

        self.real_x.append(rx)
        self.real_y.append(ry)
        self.real_z.append(rz)

        # --- 3) avg 入库 ---
        self.avg_x.append(self._last_avg['x'])
        self.avg_y.append(self._last_avg['y'])
        self.avg_z.append(self._last_avg['z'])

        # --- 4) 更新三张图 ---
        # X
        self.lx_raw.set_data(self.ticks, self.raw_x)
        self.lx_real.set_data(self.ticks, self.real_x)
        self.lx_avg.set_data(self.ticks, self.avg_x)
        self.ax_x.relim(); self.ax_x.autoscale_view()

        # Y
        self.ly_raw.set_data(self.ticks, self.raw_y)
        self.ly_real.set_data(self.ticks, self.real_y)
        self.ly_avg.set_data(self.ticks, self.avg_y)
        self.ax_y.relim(); self.ax_y.autoscale_view()

        # Z
        self.lz_raw.set_data(self.ticks, self.raw_z)
        self.lz_real.set_data(self.ticks, self.real_z)
        self.lz_avg.set_data(self.ticks, self.avg_z)
        self.ax_z.relim(); self.ax_z.autoscale_view()

        plt.pause(0.001)

def main(args=None):
    rclpy.init(args=args)
    node = PlotNode()
    try:
        while rclpy.ok():
            rclpy.spin_once(node, timeout_sec=0.05)
    except KeyboardInterrupt:
        pass
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()